package main

import (
	"adkr/pkg/adkr"
    "adkr/pkg/config"
	"adkr/pkg/crypto"
	"adkr/pkg/network"
	"encoding/json"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func main() {
	pid := flag.Int("id", -1, "Node ID")
	hostsFile := flag.String("hosts", "hosts.json", "Path to hosts file")
    paillierPub := flag.String("paillier-pub", "", "Path to Paillier public keys JSON")
    paillierPriv := flag.String("paillier-priv", "", "Path to Paillier private keys JSON (local only)")
    // Optional Threshold BLS config flags (handled elsewhere when needed)
    // _ = flag.String("bls-pub", "", "Path to Threshold BLS public config JSON")
    // _ = flag.String("bls-priv", "", "Path to Threshold BLS private config JSON (local only)")
    strongPred := flag.Bool("mvba-strong-predicate", false, "enable strong MVBA predicate (ACSS-verified dealers only)")
    abaEvidence := flag.Bool("aba-evidence", false, "attach supporter evidence in ABA messages")
    krSPBC := flag.Bool("kr-spbc", false, "use SPBC fast-path in Key Refresh (default MVBA)")
    rbcMode := flag.String("rbc-mode", "", "RBC variant: '', 'dual', or 'dual_cp'")
    coinSkipVerify := flag.Bool("coin-skip-verify", false, "skip threshold coin aggregate verification (experimental)")
    flag.Parse()
	if *pid < 0 {
		log.Fatalf("Node ID must be provided with -id flag")
	}

	jsonFile, err := os.Open(*hostsFile)
	if err != nil {
		log.Fatalf("Failed to open hosts file: %v", err)
	}
	defer jsonFile.Close()

	byteValue, _ := ioutil.ReadAll(jsonFile)
	var addrs []string
	json.Unmarshal(byteValue, &addrs)
	if *pid >= len(addrs) {
		log.Fatalf("Node ID %d is out of bounds.", *pid)
	}

	netConfig := network.NetConfig{ID: *pid, Addrs: addrs}
	n := len(addrs)
	f := (n - 1) / 3

	eSKs := make([]crypto.ZR, n)
	ePKs := make([]crypto.G1, n)
	g1Gen := crypto.GetG1()
	for i := 0; i < n; i++ {
		sk := crypto.ZRFromInt(int64(1000 + i))
		eSKs[i] = sk
		ePKs[i] = crypto.MulG1(&g1Gen, &sk)
	}
	my_eSK := eSKs[*pid]

	msgInCh := make(chan adkr.Message, 100)
	msgOutCh := make(chan adkr.Message, 100)
	net := network.NewNetwork(netConfig, msgInCh, msgOutCh)
	net.Start()
	net.WaitForFullConnection()

    // MVBA/ABA switches
    adkr.SetStrongMVBAEnabled(*strongPred)
    adkr.SetABAEvidence(*abaEvidence)
    adkr.SetKeyRefreshSPBCFast(*krSPBC)
    adkr.SetRBCMode(*rbcMode)
    adkr.SetCoinSkipVerify(*coinSkipVerify)

    // Optional: load Paillier configuration and register keys
    if *paillierPub != "" {
        pubs, err := config.LoadPaillierPublicConfig(*paillierPub)
        if err != nil { log.Fatalf("Failed to load Paillier public config: %v", err) }
        // Register all public keys
        for id, pk := range pubs { adkr.RegisterPaillierKeys(id, pk, nil) }
        // Load private only for this pid
        if *paillierPriv != "" {
            privs, err := config.LoadPaillierPrivateConfig(*paillierPriv, pubs)
            if err != nil { log.Fatalf("Failed to load Paillier private config: %v", err) }
            if sk, ok := privs[*pid]; ok {
                // re-register to include sk
                adkr.RegisterPaillierKeys(*pid, pubs[*pid], sk)
            } else {
                log.Printf("No private key for node %d in priv config", *pid)
            }
        }
    }

	// Store the return values in variables, then pass their addresses.
	g1 := crypto.GetG1()
	g2 := crypto.GetG2()
	
	// Generate threshold BLS keys for this node
	thpk, _, thsks := crypto.GenerateThresholdKeys(n, f+1)
	thpks := make([]crypto.G1, n)
	for i := 0; i < n; i++ {
		thpks[i] = crypto.MulG1(&g1, &thsks[i])
	}
	
	// Use ADKRHigh protocol for metrics output
	outputCh := make(chan adkr.ADKRHighOutput, 1)
	config := adkr.ADKRHighConfig{
		SID:   "sid-1",
		PID:   *pid,
		N:     n,
		F:     f,
		K:     1,
		G1:    &g1,
		G2:    &g2,
		EPKs:  ePKs,
		ESK:   my_eSK,
		THPK:  *thpk,
		THPKs: thpks,
		THSK:  thsks[*pid],
	}
	proto := adkr.NewADKRHighProtocol(config, msgInCh, msgOutCh, outputCh, log.Default())
	
	go func() {
		if err := proto.Run(); err != nil {
			log.Printf("ADKR-METRICS {\"node_id\":%d,\"error\":\"%v\"}", *pid, err)
			return
		}
	}()
	
	// Wait for protocol completion and output metrics
	go func() {
		select {
		case output := <-outputCh:
			// Output JSON metrics
			metrics := map[string]interface{}{
				"node_id": *pid,
				"protocol": "ADKR",
				"statistics": output.Statistics,
				"config_chain": output.ConfigChain,
			}
			jb, _ := json.Marshal(metrics)
			log.Printf("ADKR-METRICS %s", string(jb))
		case <-time.After(60 * time.Second):
			log.Printf("ADKR-METRICS {\"node_id\":%d,\"error\":\"timeout\"}", *pid)
		}
	}()

	log.Println("Node is running. Press Ctrl+C to exit.")
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	<-sigCh
}
